#!/usr/bin/env python
# encoding: utf-8
"""
@author: youfeng
@email: youfeng243@163.com
@license: Apache Licence 
@file: captcha_local.py
@time: 2016/12/5 22:11
"""
import os
import random
import time

import cv2
import numpy as np
import pytesseract
from PIL import Image
from tesserwrap import Tesseract


class CaptchaLocal(object):
    td_relative_path = '/config/tessdata'
    image_relative_path = '/image/'

    def __init__(self, log=None):
        self.log = log
        self.base_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))

    @staticmethod
    def __get_image_filename():
        return '{t}_{rand}.jpg'.format(t='%0.06f' % time.time(), rand=str(random.randint(0, 100000)))

    @staticmethod
    def __get_digit(text):
        digit_list = list(text)
        return ''.join([i for i in digit_list if '0' <= i <= '9'])

    def get_captcha_result(self, province, content):
        if province == 'sichuan':
            return self.__get_captcha_gsxt_sichuan(content)
        if province == 'court_gov':
            return self.__get_captcha_thirdpart_zx_sx(content)
        if province == 'court_gov_sx':
            return self.__get_captcha_thirdpart_zx_sx(content)
        return None

    # 获得四川验证码识别结果
    def __get_captcha_gsxt_sichuan(self, content):
        if content is None:
            return None
        if content == '':
            return None

        tmp_text = None
        filename = self.__get_image_filename()
        file_path = self.base_path + self.image_relative_path + filename
        try:
            with open(file_path, 'wb') as fp:
                fp.write(content)

            if not os.path.isfile(file_path):
                self.log.error('文件不存在: {path}'.format(path=file_path))
                return None

            img = cv2.imread(file_path)
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            _, bi = cv2.threshold(gray, 135, 255, cv2.THRESH_BINARY)
            pil_im = Image.fromarray(np.uint8(bi))
            path = self.base_path + self.td_relative_path
            tr_num1 = Tesseract(datadir=path, lang='chi_sim')
            tmp_text = tr_num1.ocr_image(pil_im)
            tmp_text = tmp_text.strip()
        except Exception as e:
            self.log.exception(e)
            tmp_text = None
        finally:
            try:
                os.remove(file_path)
            except:
                pass

        return tmp_text

    def __get_captcha_thirdpart_zx_sx(self, content):
        if content is None:
            return None
        if content == '':
            return None

        tmp_text = None
        filename = self.__get_image_filename()
        file_path = self.base_path + self.image_relative_path + filename
        try:

            with open(file_path, 'wb') as fp:
                fp.write(content)

            if not os.path.isfile(file_path):
                self.log.error('文件不存在: {path}'.format(path=file_path))
                return None

            img = Image.open(file_path)
            tmp = pytesseract.image_to_string(img)
            tmp_text = self.__get_digit(tmp)
            if not tmp_text.isdigit():
                self.log.warn('无法匹配的验证码: {tmp}'.format(tmp=tmp))
                tmp_text = None
        except Exception as e:
            self.log.exception(e)
            self.log.info(len(content))
            tmp_text = None
        finally:
            try:
                os.remove(file_path)
            except:
                pass
        return tmp_text
